import numpy as np
import matplotlib.pyplot as plt
from scipy.special import legendre
from scipy.linalg import qr, solve

# 定义Legendre多项式作为基函数
def basis_functions(x, m):
    """生成m个Legendre多项式作为基函数"""
    return np.array([legendre(k)(x) for k in range(m)]).T

# 定义DE约束嵌入的表达式
def constrained_expression(x, g, Y0, Yf, p, q, eta1, eta2):
    """构建约束表达式 y(x)"""
    alpha = eta1 * p + eta2 * q
    return g - alpha * Y0 - (1 - alpha) * Yf

# 最小二乘法求解DE
def least_squares_solution(f2, f1, f0, f, x, Y0, Yf, m):
    """
    使用最小二乘法求解线性微分方程。
    
    参数:
        f2, f1, f0: DE系数函数
        f: 非齐次项函数
        x: 离散点数组
        Y0, Yf: 初始条件和边界条件
        m: 基函数数量
    返回:
        y: 解函数值
    """
    # 生成基函数及其导数
    h = basis_functions(x, m)
    h_prime = np.gradient(h, x, axis=0, edge_order=2)
    h_double_prime = np.gradient(h_prime, x, axis=0, edge_order=2)

    # 构造p(x)和q(x)
    p = 1 - x
    q = x

    # 计算η1和η2
    M = np.array([[p[0], q[0]], [p[-1], q[-1]]])
    v = np.array([Y0, Yf])
    eta = np.linalg.solve(M, v)

    # 构造约束表达式
    g = np.dot(h, np.ones(m))  # 初始化g为自由函数
    alpha = eta[0] * p + eta[1] * q
    y = g - alpha * Y0 - (1 - alpha) * Yf

    # 构造最小二乘矩阵A
    A = (f2[:, None] * h_double_prime +
         f1[:, None] * h_prime +
         f0[:, None] * h)
    b = f - (f2 * alpha +
             f1 * alpha +
             f0 * alpha)

    # QR分解求解最小二乘问题
    Q, R = qr(A, mode='economic')
    xi = solve(R, np.dot(Q.T, b))

    # 计算最终解
    g_final = np.dot(h, xi)
    y_final = g_final - alpha * Y0 - (1 - alpha) * Yf
    return y_final

# 示例：求解一个BVP
if __name__ == "__main__":
    # 定义DE系数函数
    def f2(x): return 1 * np.ones_like(x)
    def f1(x): return 2 * np.ones_like(x)
    def f0(x): return 1 * np.ones_like(x)
    def f(x): return np.zeros_like(x)

    # 边界条件
    Y0 = 1
    Yf = 3

    # 离散化变量
    x = np.linspace(-1, 1, 100)

    # 基函数数量
    m = 14

    # 求解
    y_solution = least_squares_solution(f2(x), f1(x), f0(x), f(x), x, Y0, Yf, m)

    # 绘图
    plt.plot(x, y_solution, label="Numerical Solution")
    plt.xlabel("x")
    plt.ylabel("y(x)")
    plt.title("Least-Squares Solution of BVP")
    plt.legend()
    plt.show()
